import math
import os
from dataprocess.datareprepare import prepare_data_loaders
from model.dborgotmodel import dborgotmodel
from model.dbmodel import dbmodel
from model.gotmodel import gotmdmodel
from utils.trainresultdraw import plot_accuracy_and_loss
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
from utils.dborgottrainprocess import dborgottrain
from utils.dbtrainprocess import dbtrain
from utils.gotmdtrainprocess import gotmdtrain

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
device = "cuda"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# cuda
model = dborgotmodel().to(device)
dbmodel = dbmodel().to(device)
gotmdmodel = gotmdmodel().to(device)
train_loader, test_loader = prepare_data_loaders()
# 损失函数，交叉熵
Loss_fn = nn.CrossEntropyLoss().to(device)
Yh_fn = torch.optim.SGD(model.parameters(), lr=5e-4, momentum=0.9)

db_train_loss_list = []
gotmd_train_loss_list = []
db_train_acc_list = []
gotmd_train_acc_list = []

# 糖尿病分期模型的优化器
db_Yh_fn = torch.optim.SGD(dbmodel.parameters(), lr=5e-4, momentum=0.9)

# 痛风分期模型的优化器
gotmd_Yh_fn = torch.optim.SGD(gotmdmodel.parameters(), lr=5e-4, momentum=0.9)

t_a, t_loss, all_x_db, all_y_db, all_x_gotmd, all_y_gotmd = dborgottrain(train_loader, model, Loss_fn, Yh_fn)

print('疾病分类预测正确率:' + str(t_a) + '，训练Loss:' + str(t_loss))
all_x_db_tensor = torch.cat(all_x_db, dim=0)
all_y_db_tensor = torch.cat(all_y_db, dim=0)
all_x_gotmd_tensor = torch.cat(all_x_gotmd, dim=0)
all_y_gotmd_tensor = torch.cat(all_y_gotmd, dim=0)
x_db_length = all_x_db_tensor.size(0)
x_gotmd_length = all_x_gotmd_tensor.size(0)
nearest_multiple_128 = int(128 * math.floor(x_db_length / 128.0))
gotmd_nearest_multiple_128 = int(128 * math.floor(x_gotmd_length / 128.0))
all_x_db_tensor = all_x_db_tensor[:nearest_multiple_128]
all_y_db_tensor = all_y_db_tensor[:nearest_multiple_128]
all_x_gotmd_tensor = all_x_db_tensor[:gotmd_nearest_multiple_128]
all_y_gotmd_tensor = all_y_db_tensor[:gotmd_nearest_multiple_128]
print(all_x_db_tensor.shape)
print(all_x_gotmd_tensor.shape)
gotmd_dataset = TensorDataset(all_x_gotmd_tensor, all_y_gotmd_tensor)
db_dataset = TensorDataset(all_x_db_tensor, all_y_db_tensor)
db_dataset = DataLoader(dataset=db_dataset, batch_size=128, shuffle=True)
gotmd_dataset = DataLoader(dataset=gotmd_dataset, batch_size=128, shuffle=True)
db_epoch = 50
for t in range(db_epoch):
    print(f'批次{t + 1}训练')
    db_a, db_loss = dbtrain(db_dataset, dbmodel, Loss_fn, db_Yh_fn, t)
    gotmd_a, gotmd_loss = gotmdtrain(gotmd_dataset, gotmdmodel, Loss_fn, gotmd_Yh_fn, t)
    db_train_loss_list.append(db_loss)
    db_train_acc_list.append(db_a)
    gotmd_train_loss_list.append(gotmd_loss)
    gotmd_train_acc_list.append((gotmd_a))
    print('糖尿病训练正确率:' + str(db_a) + '，训练Loss:' + str(db_loss))
    print('痛风训练正确率:' + str(gotmd_a) + '，训练Loss:' + str(gotmd_loss))
plot_accuracy_and_loss(db_train_acc_list, gotmd_train_acc_list, db_train_loss_list, gotmd_train_loss_list)
print('finish')
